#!/usr/bin/python
#Author: Jimmy Saw
#Usage: python get_lineage.py "Genus species"
#It is supposed to give you the parent phylum of the query term

import sys
import re
from Bio import Entrez

Entrez.email = "jimmysaw@gmail.com"

filename = sys.argv[1]
sf = open(filename, "rU")
sfl = sf.readlines()

rx1 = re.compile('(\w+ \w+).*')
rx2 = re.compile('(\w+ \w+ \w+).*')
rx3 = re.compile('(\w+).*') #keep only the genus name as search query
rx4 = re.compile('candidatus', re.IGNORECASE)

for line in sfl:
    searchterm = line.rstrip()
    if rx4.match(searchterm):
        pattern = rx2.match(searchterm)
        handle = Entrez.esearch(db="Taxonomy", term=pattern.group(1))
    else:
        pattern = rx3.match(searchterm)
        handle = Entrez.esearch(db="Taxonomy", term=pattern.group(1))

    record = Entrez.read(handle)
    taxid = record["IdList"]

    if len(taxid) != 0:

        h = Entrez.efetch(db="Taxonomy", id=taxid, retmode="xml")

        rec = Entrez.parse(h)

        for r in rec:
            lineage = r.get("Lineage")
            l = lineage.split(';')
            #print l[3].lstrip()
            #print "Search term used =", searchterm #pattern.group(1)
            #print "Lineage found =", lineage
            #print "Phylum =", l[2].lstrip()
            print searchterm + "\t" + l[3].lstrip()
    else:
        print "Can't query NCBI Taxonomy. Taxid not found!"
        print "Search term used =", searchterm

sf.close()